package com.small.xx.ai.spring.rag.controller;

import com.small.xx.ai.spring.rag.processor.DocumentSelectFirst;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/rag/simple")
public class SimpleRagController {

    private static final Logger logger = LoggerFactory.getLogger(SimpleRagController.class);
    private final SimpleVectorStore simpleVectorStore;
    private final ChatClient chatClient;
    private final ChatClient.Builder chatClientBuilder;

    public SimpleRagController(EmbeddingModel embeddingModel, ChatClient.Builder builder) {
        this.simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();
        this.chatClientBuilder = builder;
        this.chatClient = builder.build();
    }

    @GetMapping("/add")
    public void add() {
        logger.info("start add data");
        HashMap<String, Object> map = new HashMap<>();
        map.put("year", 2025);
        map.put("name", "xiaoxiao");
        List<Document> documents = List.of(new Document("你的姓名是小小，四川本地人，是一位自由工作者"),
            new Document("你的姓名是小小，专业领域包含的打游戏，玩儿", Map.of("year", 2024)),
            new Document("你姓名是小小，爱好是发呆、思考、运动", map));
        simpleVectorStore.add(documents);
    }

    @GetMapping("/chat")
    public String chat(
        @RequestParam(value = "query", defaultValue = "你好，请告诉我小小这个人的身份信息") String query) {
        logger.info("start chat");
        return chatClient.prompt(query).call().content();
    }

    @GetMapping("/chat-rag-advisor")
    public String chatRagAdvisor(
        @RequestParam(value = "query", defaultValue = "你好，请告诉我小小这个人的身份信息") String query) {
        logger.info("start chat with rag-advisor");
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
            .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(simpleVectorStore).build())
            // 上下文查询增强器（ContextualQueryAugmenter），用于在检索增强生成（RAG）流程中，将检索到的文档上下文与原始查询结合，提升大模型回答的相关性和准确性
            .queryAugmenter(ContextualQueryAugmenter.builder()
                // 当检索流程未找到任何相关文档（即上下文为空）时，依然允许增强器正常工作（不会因缺少上下文而报错）。
                //此时，增强器会直接返回原始查询（或仅做轻微调整），确保大模型能基于自身知识继续生成回答。
                .allowEmptyContext(true)
                //完成配置并创建增强器实例。
                .build())
            //
            .build();

        return chatClient.prompt(query).advisors(retrievalAugmentationAdvisor).call().content();
    }

    /**
     * <pre>
     * 这段代码主要围绕检索增强生成（RAG）流程展开，除了显式定义的类之外，还隐含依赖了多个工具类（通常来自同框架的底层模块或关联组件）。以下是核心工具类及作用分析：
     * 1. 大语言模型客户端相关工具类
     * ChatClient由 this.chatClientBuilder.build() 生成的实例，是与大语言模型交互的核心工具。负责执行实际的文本生成任务（如查询扩展、翻译、上下文增强等），是 MultiQueryExpander、TranslationQueryTransformer 等组件的底层依赖。
     * ChatClientBuilder用于构建 ChatClient 的建造者类，封装了模型连接参数（如 API 密钥、模型名称、请求超时时间等），代码中通过 this.chatClientBuilder 传递给多个组件，作为模型交互的入口。
     * 2. 向量存储相关工具类
     * VectorStore（以 simpleVectorStore 为例）向量数据库的抽象接口或实现类（如常见的 Chroma、Weaviate、Milvus 等），负责存储文档的向量表示，并支持基于向量相似度的检索。VectorStoreDocumentRetriever 通过它执行实际的文档检索操作。
     * Document文档数据的封装类，用于存储从向量库中检索到的文本内容及元数据（如页码、来源等）。是 VectorStoreDocumentRetriever 的返回类型，也是 ConcatenationDocumentJoiner、DocumentSelectFirst 等组件的处理对象。
     * 3. 文本处理与转换工具类
     * Query封装用户原始查询的类，包含查询文本及可能的附加信息（如语言类型、上下文标识等），是 queryExpander、queryTransformers、queryAugmenter 等组件的输入参数。
     * StringJoiner 或类似文本拼接工具ConcatenationDocumentJoiner 内部可能依赖此类，用于将多个 Document 的文本内容按规则拼接为一个连续字符串（如添加分隔符、去重等）。
     * 4. 流程编排与抽象工具类
     * AdvisorRetrievalAugmentationAdvisor 的父接口或抽象类，定义了 “增强器” 的通用行为（如接收查询、处理并返回增强结果），使不同类型的增强逻辑（如 RAG、工具调用等）可统一接入大语言模型的调用流程。
     * Builder 模式相关基类所有组件（如 MultiQueryExpander.builder()、RetrievalAugmentationAdvisor.builder()）使用的建造者基类，提供统一的对象构建接口（如 build() 方法、属性设置方法），简化复杂对象的初始化。
     * 总结
     * 这些工具类共同支撑了 RAG 流程的核心能力：
     * 大语言模型交互（ChatClient 系列）
     * 向量数据存储与检索（VectorStore 系列）
     * 文本与查询的结构化处理（Document、Query 等）
     * 组件的灵活组装（Builder 模式、Advisor 接口）
     * </pre>
     *
     * @param query
     * @return
     */
    @GetMapping("/chat-rag-advisor-detail")
    public String chatRagAdvisorByDetail(
        @RequestParam(value = "query", defaultValue = "你好，请告诉我小小这个人的身份信息") String query) {
        logger.info("start chat with rag-advisor");

        // 1. 检索前处理阶段（Pre-Retrieval）：对原始查询进行扩展和转换，提升检索召回率
        // 1.1 初始化多查询扩展器（MultiQueryExpander）
        // 作用：通过大语言模型将原始查询扩展为多个语义相似的变体查询，增加检索覆盖范围
        MultiQueryExpander multiQueryExpander = MultiQueryExpander.builder()  // 使用建造者模式构建实例
            .chatClientBuilder(this.chatClientBuilder)  // 传入大语言模型客户端构建器（用于生成扩展查询）
            .build();  // 完成构建

        // 1.2 初始化查询翻译转换器（TranslationQueryTransformer）
        // 作用：将原始查询翻译为目标语言，适配多语言向量库的检索需求
        TranslationQueryTransformer translationQueryTransformer = TranslationQueryTransformer.builder()  // 使用建造者模式构建实例
            .chatClientBuilder(this.chatClientBuilder)  // 传入大语言模型客户端构建器（用于执行翻译）
            .targetLanguage("English")  // 设置目标翻译语言为英文（若向量库内容为英文时提升匹配度）
            //                .targetLanguage("Chinese")  // 可选：若向量库为中文，可切换为中文
            .build();  // 完成构建

        // 2. 检索阶段（Retrieval）：根据处理后的查询从向量库中获取相关文档
        // 2.1 初始化向量存储文档检索器（VectorStoreDocumentRetriever）
        // 作用：连接向量数据库，根据查询向量检索最相似的文档
        VectorStoreDocumentRetriever vectorStoreDocumentRetriever =
            VectorStoreDocumentRetriever.builder()  // 使用建造者模式构建实例
                .vectorStore(simpleVectorStore)  // 传入具体的向量存储实例（存储文档向量的数据库）
                .build();  // 完成构建

        // 2.2 初始化文档拼接器（ConcatenationDocumentJoiner）
        // 作用：将检索到的多个文档内容按顺序拼接为一个字符串，便于后续作为上下文输入大模型
        ConcatenationDocumentJoiner concatenationDocumentJoiner = new ConcatenationDocumentJoiner();  // 直接实例化（无额外配置）

        // 3. 检索后处理阶段（Post-Retrieval）：对检索到的文档进行过滤或筛选，优化上下文质量
        // 3.1 初始化文档选择器（DocumentSelectFirst）
        // 作用：从检索到的多个文档中仅保留第一个（通常是相似度最高的文档）
        DocumentSelectFirst documentSelectFirst = new DocumentSelectFirst();  // 直接实例化（无额外配置）

        // 4. 生成阶段（Generation）：结合检索到的上下文增强查询，提升大模型回答准确性
        // 4.1 初始化上下文查询增强器（ContextualQueryAugmenter）
        // 作用：将检索到的文档上下文注入原始查询，让大模型生成回答时更贴合上下文信息
        ContextualQueryAugmenter contextualQueryAugmenter = ContextualQueryAugmenter.builder()  // 使用建造者模式构建实例
            .allowEmptyContext(true)  // 允许上下文为空（当检索不到文档时仍能生成回答）
            .build();  // 完成构建

        // 组装检索增强顾问（RetrievalAugmentationAdvisor）：整合上述所有组件，形成完整的RAG（检索增强生成）流程
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor =
            RetrievalAugmentationAdvisor.builder().queryExpander(multiQueryExpander)  // 绑定多查询扩展器（将原始查询扩展为3个变体，默认配置）
                .queryTransformers(translationQueryTransformer)  // 绑定查询翻译器（将扩展后的查询转为英文）
                .documentRetriever(vectorStoreDocumentRetriever)  // 绑定向量检索器（根据处理后的查询检索文档）
                .documentJoiner(concatenationDocumentJoiner)  // 绑定文档拼接器（将检索到的文档合并为上下文）
                .documentPostProcessors(documentSelectFirst)  // 绑定文档后处理器（筛选出最相关的第一个文档）
                .queryAugmenter(contextualQueryAugmenter)  // 绑定上下文增强器（用筛选后的文档增强查询）
                .build();  // 完成RAG流程组装

        // 执行完整流程并返回结果：
        // 1. 构建大语言模型客户端 -> 2. 传入原始查询 -> 3. 应用RAG增强顾问 -> 4. 调用模型生成回答 -> 5. 返回回答内容
        return this.chatClientBuilder.build()  // 构建大语言模型客户端
            .prompt(query)  // 设置用户原始查询
            .advisors(retrievalAugmentationAdvisor)  // 应用检索增强逻辑
            .call()  // 调用模型生成回答
            .content();  // 提取并返回回答内容
    }
}
